"""
The super twisting algorithm should also work for extra damping terms

"""


from scipy.integrate import solve_ivp
import numpy as np
import matplotlib.pyplot as plt
from dataclasses import dataclass

@dataclass
class Params:
    k1:float
    k2:float
    def delta(self,t):
        return 0.1*np.sin(8*t)

def fun(t,x,p:Params):
    x1=x[0]
    x2=x[1]
    signals={
        "delta":p.delta(t),
        "x1p":abs(x1)**(1/2)*np.sign(x1)
    }
    dx1=x2-p.k1*signals["x1p"]
    dx2=-p.k2*np.sign(x1)+signals["delta"]-x1-x2
    return np.array([dx1,dx2]),signals

def rhs(t,x,p:Params):
    return fun(t,x,p)[0]
    
def main():
    x0=np.array([1,1])
    p=Params(1,1)
    sol=solve_ivp(rhs,[0,60],x0,args=(p,),t_eval=np.linspace(0,10,1000),method="BDF")

    debug=[fun(t,sol.y[:,i],p)[1] for i,t in enumerate(sol.t)]

    plt.figure()
    plt.subplot(2,1,1)
    plt.title("Stabilization of x1")
    plt.plot(sol.t,sol.y[0],label="x1")
    plt.subplot(2,1,2)
    plt.plot(sol.t,[s["delta"] for s in debug],label="delta")
    plt.plot(sol.t,[s["x1p"] for s in debug],label=r"$|x_1|^{1/2}\mathrm{sign}(x_1)$")
    plt.plot(sol.t,sol.y[1],label="x2")
    plt.title(f"x1={sol.y[0][-1]:e},x2={sol.y[1][-1]:e}")
    plt.legend()
    plt.tight_layout()
    plt.show()
    return sol

if __name__=="__main__":
    main()